#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
聚类检索器功能测试

测试DBSCANClusteringRetriever和HierarchicalClusteringRetriever的高级功能
"""

import sys
import os

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 在导入任何库之前强制设置环境变量
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_CO_URL_HOME"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.expanduser("~/.cache/huggingface")
os.environ["SENTENCE_TRANSFORMERS_HOME"] = os.path.expanduser("~/.cache/sentence_transformers")

# 强制导入配置，确保镜像环境变量已设置
import questionretrieval.config
from loguru import logger

logger.info("🌐 使用 HuggingFace 国内镜像配置")
logger.info(f"📁 缓存目录: {questionretrieval.config.HF_HUB_CACHE}")
logger.info(f"🔗 镜像地址: {questionretrieval.config.HF_MIRROR_ENDPOINT}")

import pytest
import numpy as np
from questionretrieval import load_questions

# 测试数据
TEST_QUESTIONS = [
    "大伙房水库的主管部门是哪个单位？",
    "大伙房水库总库容/防洪库容/兴利库容/死库容？",
    "大伙房水库调度目标和任务是什么？",
    "大伙房水库防洪调度的核心任务是什么？",
    "大伙房水库供水调度的核心任务是什么？",
    "水库的防洪标准是什么？",
    "如何进行水库调度？",
    "水库管理的基本原则",
    "水库安全监测要求",
    "水库运行维护标准"
]

TEST_QUERIES = [
    "大伙房水库的管理和调度",
    "水库的容量信息",
    "防洪相关任务"
]

def get_test_questions():
    """
    获取测试问题
    """
    try:
        questions = load_questions()
        if len(questions) > 0:
            return questions
    except:
        pass
    return TEST_QUESTIONS

def test_dbscan_clustering_parameters():
    """
    测试 DBSCAN 聚类检索器参数调优功能
    """
    logger.info("\n=== 测试 DBSCAN 聚类检索器参数调优 ===")
    
    try:
        from questionretrieval import DBSCANClusteringRetriever
        questions = get_test_questions()
        
        # 测试不同参数组合
        parameter_sets = [
            {'eps': 0.3, 'min_samples': 2},
            {'eps': 0.5, 'min_samples': 3},
            {'eps': 0.7, 'min_samples': 2}
        ]
        
        results = {}
        
        for i, params in enumerate(parameter_sets):
            try:
                retriever = DBSCANClusteringRetriever(eps=params['eps'], min_samples=params['min_samples'])
                retriever.load_candidates(questions)
                
                # 获取聚类信息
                cluster_info = retriever.get_cluster_info()
                results[f"参数组{i+1}"] = {
                    'params': params,
                    'clusters': cluster_info.get('n_clusters', 0) if cluster_info else 0,
                    'noise_points': cluster_info.get('n_noise', 0) if cluster_info else 0,
                    'success': True
                }
                
                logger.info(f"  参数组{i+1} (eps={params['eps']}, min_samples={params['min_samples']}):")                
                logger.info(f"    聚类数量: {results[f'参数组{i+1}']['clusters']}")                
                logger.info(f"    噪声点数: {results[f'参数组{i+1}']['noise_points']}")
                
            except Exception as e:
                results[f"参数组{i+1}"] = {
                    'params': params,
                    'success': False,
                    'error': str(e)
                }
                logger.error(f"  参数组{i+1} 失败: {e}")
        
        success_count = sum(1 for r in results.values() if r['success'])
        
        if success_count > 0:
            logger.info(f"✓ 参数调优测试成功，{success_count}/{len(parameter_sets)} 个参数组有效")
            return True
        else:
            logger.warning("✗ 所有参数组测试失败")
            return False
    
    except ImportError:
        logger.warning("✗ DBSCANClusteringRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ DBSCAN参数调优测试失败: {e}")
        return False

def test_dbscan_outlier_detection():
    """
    测试 DBSCAN 聚类检索器异常值检测功能
    """
    logger.info("\n=== 测试 DBSCAN 聚类检索器异常值检测 ===")
    
    try:
        from questionretrieval import DBSCANClusteringRetriever
        retriever = DBSCANClusteringRetriever(eps=0.5, min_samples=2)
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        
        # 获取聚类信息来检测异常值（噪声点）
        cluster_info = retriever.get_cluster_info()
        
        if cluster_info is not None:
            noise_count = cluster_info.get('n_noise', 0)
            logger.info(f"✓ 异常值检测成功，发现 {noise_count} 个噪声点")
            return True
        else:
            logger.warning("⚠️  无法获取聚类信息")
            return False
    
    except ImportError:
        logger.warning("✗ DBSCANClusteringRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 异常值检测测试失败: {e}")
        return False

def test_hierarchical_clustering_parameters():
    """
    测试层次聚类检索器参数调优功能
    """
    logger.info("\n=== 测试层次聚类检索器参数调优 ===")
    
    try:
        from questionretrieval import HierarchicalClusteringRetriever
        questions = get_test_questions()
        
        # 测试不同参数组合（HierarchicalClusteringRetriever接受linkage_method参数）
        parameter_sets = [
            {'linkage_method': 'ward'},
            {'linkage_method': 'complete'},
            {'linkage_method': 'average'}
        ]
        
        results = {}
        
        for i, params in enumerate(parameter_sets):
            try:
                retriever = HierarchicalClusteringRetriever(
                    linkage_method=params['linkage_method']
                )
                retriever.load_candidates(questions)
                
                # 获取聚类信息
                cluster_info = retriever.get_cluster_info()
                results[f"参数组{i+1}"] = {
                    'params': params,
                    'clusters': cluster_info.get('n_clusters', 0) if cluster_info else 0,
                    'success': True
                }
                
                logger.info(f"  参数组{i+1} (linkage_method={params['linkage_method']}):")
                logger.info(f"    聚类数量: {results[f'参数组{i+1}']['clusters']}")
                
            except Exception as e:
                results[f"参数组{i+1}"] = {
                    'params': params,
                    'success': False,
                    'error': str(e)
                }
                logger.error(f"  参数组{i+1} 失败: {e}")
        
        success_count = sum(1 for r in results.values() if r['success'])
        
        if success_count > 0:
            logger.info(f"✓ 参数调优测试成功，{success_count}/{len(parameter_sets)} 个参数组有效")
            return True
        else:
            logger.warning("✗ 所有参数组测试失败")
            return False
    
    except ImportError:
        logger.warning("✗ HierarchicalClusteringRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 层次聚类参数调优测试失败: {e}")
        return False

def test_hierarchical_clustering_dendrogram():
    """
    测试层次聚类检索器树状图功能
    """
    logger.info("\n=== 测试层次聚类检索器树状图功能 ===")
    
    try:
        from questionretrieval import HierarchicalClusteringRetriever
        import matplotlib
        matplotlib.use('Agg')  # 使用非交互式后端
        import matplotlib.pyplot as plt
        
        # 设置中文字体支持
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
        
        retriever = HierarchicalClusteringRetriever()
        questions = get_test_questions()[:6]  # 使用较少问题
        
        retriever.load_candidates(questions)
        
        # 尝试生成树状图
        try:
            import os
            reports_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "reports")
            os.makedirs(reports_dir, exist_ok=True)
            dendrogram_path = os.path.join(reports_dir, "test_dendrogram.png")
            retriever.plot_dendrogram(save_path=dendrogram_path, show_labels=True)
            logger.info(f"✓ 树状图生成成功，已保存为 {dendrogram_path}")
            
            # 检查聚类信息
            cluster_info = retriever.get_cluster_info()
            if cluster_info:
                levels = cluster_info.get('levels', [])
                if isinstance(levels, list):
                    logger.info(f"  聚类层级数: {len(levels)}")
                else:
                    logger.info(f"  聚类层级数: {levels}")
                logger.info(f"  总聚类数: {cluster_info.get('total_clusters', 0)}")
            
            return True
        except Exception as plot_error:
            logger.warning(f"⚠️  树状图生成失败: {plot_error}")
            # 检查基本聚类功能
            cluster_info = retriever.get_cluster_info()
            if cluster_info is not None:
                logger.info("✓ 层次聚类基本功能正常")
            return True
        else:
            logger.warning("✗ 层次聚类功能异常")
            return False
    
    except ImportError:
        logger.warning("✗ HierarchicalClusteringRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 层次聚类功能测试失败: {e}")
        return False

def test_clustering_batch_processing():
    """
    测试聚类检索器批量处理功能
    """
    logger.info("\n=== 测试聚类检索器批量处理 ===")
    
    try:
        from questionretrieval import DBSCANClusteringRetriever
        retriever = DBSCANClusteringRetriever()
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        
        # 测试多个查询的推荐
        batch_results = {}
        for query in TEST_QUERIES:
            try:
                results = retriever.recommend(query, top_k=2)
                batch_results[query] = results
            except Exception as e:
                logger.error(f"  查询失败: {query} - {e}")
        
        if batch_results:
            logger.info(f"✓ 批量处理成功，处理了 {len(batch_results)} 个查询")
            for query, results in batch_results.items():
                logger.info(f"  查询: {query}")
                for i, result in enumerate(results[:2], 1):
                    if isinstance(result, tuple) and len(result) >= 3:
                        question, score, similarity = result[0], result[1], result[2]
                        logger.info(f"    {i}. {question} (分数: {score:.3f}, 相似度: {similarity:.3f})")
                    elif isinstance(result, dict):
                        logger.info(f"    {i}. {result.get('question', 'N/A')} (相似度: {result.get('similarity', 0):.3f})")
                    else:
                        logger.info(f"    {i}. {result}")
            return True
        else:
            logger.warning("✗ 批量处理失败，无结果")
            return False
            
    except ImportError:
        logger.warning("✗ DBSCANClusteringRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 批量处理测试失败: {e}")
        return False

def test_all_clustering_features():
    """
    测试所有聚类检索器功能
    """
    logger.info("\n" + "="*60)
    logger.info("聚类检索器功能测试")
    logger.info("="*60)
    
    results = {
        'dbscan_parameters': test_dbscan_clustering_parameters(),
        'dbscan_outliers': test_dbscan_outlier_detection(),
        'hierarchical_parameters': test_hierarchical_clustering_parameters(),
        'hierarchical_dendrogram': test_hierarchical_clustering_dendrogram(),
        'clustering_batch': test_clustering_batch_processing()
    }
    
    logger.info("\n" + "-"*50)
    logger.info("聚类功能测试结果")
    logger.info("-"*50)
    
    success_count = 0
    for name, success in results.items():
        status = "✓ 成功" if success else "✗ 失败"
        logger.info(f"{name:25} : {status}")
        if success:
            success_count += 1
    
    logger.info(f"\n总计: {success_count}/{len(results)} 个测试通过")
    
    if success_count > 0:
        logger.info("\n💡 提示: 聚类检索器支持多种高级功能")
        logger.info("   - 参数调优 (eps, min_samples, n_clusters, linkage)")
        logger.info("   - 异常值检测")
        logger.info("   - 树状图生成")
        logger.info("   - 批量处理")
    else:
        logger.warning("\n⚠️  所有聚类功能测试失败，请检查依赖安装")
        logger.info("解决方案: pip install scikit-learn")
    
    return results

if __name__ == "__main__":
    test_all_clustering_features()